# %%
import torch
import torch.nn as nn
# %%
class BasicExpert(nn.Module):
    def __init__(self, feature_in,  feature_out):
        super().__init__()
        self.linear = nn.Linear(feature_in, feature_out)

    def forward(self, x):
        return self.linear(x)
# %%
class BasicMOE(nn.Module):
    def __init__(self, feature_in, feature_out, expert_number) -> None:
        super().__init__()
        self.experts = nn.ModuleList(
            [
                BasicExpert(feature_in, feature_out) for _ in range(expert_number)
            ]
        )

        self.gate = nn.Linear(feature_in, expert_number)

    def forward(self, x):
        expert_weights = self.gate(x)

        expert_out_list = [
            expert(x).unsqueeze(1) for expert in self.experts
        ]

        expert_out = torch.cat(expert_out_list, dim=1)

        expert_weight = expert_weights.unsqueeze(1)

        output = expert_weight @ expert_out

        return output
# %%
def test_basic_mode():
    x = torch.rand(2, 4)

    basic_moe = BasicMOE(4,3,3)
    out = basic_moe(x)
    print(out)

# %%
test_basic_mode()